from numpy import (
    ndarray,
    mean,
    std,
)

from src.data.precoder.mmse_precoder import (
    mmse_precoder_no_norm,
)
from src.utils.real_complex_vector_reshaping import (
    complex_vector_to_double_real_vector,
)


def add_random_distribution(
        rng,
        action: ndarray,
        tau_momentum: float,
) -> ndarray:
    """
    将一个动作向量与同长度的随机分布向量混合，
    混合方式为 tau_momentum * random_distribution + (1 - tau_momentum) * action。

    参数:
    rng: 随机数生成器
    action: 输入的动作向量，类型为ndarray
    tau_momentum: 混合系数，控制随机分布和原始动作的加权比例，范围为[0, 1]

    返回:
    noisy_action: 混合后的动作向量，类型为ndarray
    """

    # 如果 tau_momentum 为0，直接返回原始动作
    if tau_momentum == 0.0:
        return action

    # 生成随机分布向量
    # 使用正态分布生成随机向量，其均值为action的均值，标准差为action的标准差
    random_distribution = rng.normal(loc=mean(action), scale=std(action))

    # 按比例混合随机分布和原始动作向量
    noisy_action = tau_momentum * random_distribution + (1 - tau_momentum) * action

    # 返回混合后的动作向量
    return noisy_action


def add_mmse_precoder(
        action: ndarray,
        tau_momentum: float,
        channel_matrix,
        noise_power_watt,
        power_constraint_watt,
) -> ndarray:
    """
    使用MMSE预编码器将动作向量进行混合。
    混合方式为 tau_momentum * mmse_precoder + (1 - tau_momentum) * action。

    参数:
    action: 输入的动作向量，类型为ndarray
    tau_momentum: 混合系数，控制MMSE预编码器输出和原始动作的加权比例，范围为[0, 1]
    channel_matrix: 信道矩阵
    noise_power_watt: 噪声功率，单位为瓦特
    power_constraint_watt: 功率约束，单位为瓦特

    返回:
    noisy_action: 混合后的动作向量，类型为ndarray
    """

    # 如果 tau_momentum 为0，直接返回原始动作
    if tau_momentum == 0.0:
        return action

    # 计算MMSE预编码器输出
    w_mmse = mmse_precoder_no_norm(
        channel_matrix=channel_matrix,
        noise_power_watt=noise_power_watt,
        power_constraint_watt=power_constraint_watt,
    )
    
    # 将预编码器输出展平
    w_mmse = w_mmse.flatten()
    
    # 将复数向量转换为双实数向量
    w_mmse_real = complex_vector_to_double_real_vector(w_mmse)

    # 按比例混合MMSE预编码器输出和原始动作向量
    noisy_action = tau_momentum * w_mmse_real + (1 - tau_momentum) * action

    # 返回混合后的动作向量
    return noisy_action
